import os,sys
import re
import string
import warnings
import string
import collections
import math
import sets
import random


outfile="qualityplots"
output = outfile + ".qual.R"
FO=open(output,'w')
shrink=1000
read_len=0
quality = {}
q_max = -1
q_min = 10000
q_list=[]
i_box={}
try:

   file=sys.argv[1] #out.sam or bowtie out file. if it is sam format extension has to be sam
except:
   print "Please enter input file: sam or bowtie out file or fastq file. if it is sam format extension has to be sam and fastq file extension has tpo be faq"
   sys.exit(2)
try:
  fp=open(sys.argv[1],"r")
except:
  print "File "+sys.argv[1]+" not found"
  sys.exit(2)

if sys.argv[1].split(".")[1]=="sam":
  

   for it in fp:
     if it.count("@")==0:
        its=it.strip().split("\t")
        if its[2].count("chr")>0:
           qual_str=its[10]
           if its[1]!="0":
              tmp=list(its[10])
              tmp.reverse()
              qual_str="".join(tmp)
        
           if read_len<len(qual_str):
              read_len=len(qual_str)
           for i,j in enumerate(qual_str):
             q=ord(j)-33
              
             if q > q_max: q_max = q
             if q < q_min: q_min = q
             if quality.has_key(i)==0:
                 quality[i]={}
             if quality[i].has_key(q)==0:
                   quality[i][q] = 0
             quality[i][q] += 1

elif sys.argv[1].split(".")[1]=="fastq":
   ln=fp.readline()
   while ln:
       ln1=ln
       ln2=fp.readline()
       ln3=fp.readline()
       qual_str=fp.readline().strip()
       if read_len<len(qual_str):
              read_len=len(qual_str)
       for i,j in enumerate(qual_str):
             q=ord(j)-33
             
             if q > q_max: q_max = q
             if q < q_min: q_min = q
             if quality.has_key(i)==0:
                 quality[i]={}
             if quality[i].has_key(q)==0:
                   quality[i][q] = 0
             quality[i][q] += 1
       ln=fp.readline()
      
   
   
else:
   for it in fp:
     
        its=it.strip().split("\t")
        if its[2].count("chr")>0:
           qual_str=its[5]
           if its[1]!="+":
              tmp=list(its[5])
              tmp.reverse()
              qual_str="".join(tmp)
        
           if read_len<len(qual_str):
              read_len=len(qual_str)
           for i,j in enumerate(qual_str):
             q=ord(j)-33
             if q > q_max: q_max = q
             if q < q_min: q_min = q
             if quality.has_key(i)==0:
                 quality[i]={}
             if quality[i].has_key(q)==0:
                   quality[i][q] = 0
             quality[i][q] += 1

					
fp.close()
		
for p in range(0,read_len):
   val=[]
   occurrence=[]
   for q in range(q_min,q_max+1):
      if quality.has_key(p) and quality[p].has_key(q):
         val.append(str(q))
         occurrence.append(str(quality[p][q]))
         q_list.append(str(quality[p][q]))
      else:
         q_list.append(str(0))
   i_box[p] = 'rep(c(' + ','.join(val) + '),times=c(' + ','.join(occurrence) + '))'
		
		
#generate R script for boxplot
print >>FO, "pdf(\'%s\')" % (outfile + ".qual.boxplot.pdf")
for i in sorted(i_box):
    print >>FO,'p'+str(i) + '<-' + i_box[i]
    
print >>FO, 'boxplot(' + ','.join(['p'+str(i) for i in i_box]) + ',xlab=\"Position of Read(5\'->3\')\",ylab=\"Phred Quality Score\",outline=F' + ', ylim = c('+str(q_min)+','+str(q_max)+ '))'
print >>FO,"dev.off()"
		
		
#generate R script for heatmap
print >>FO, '\n'
print >>FO, "pdf(\'%s\')" % (outfile + ".qual.heatmap.pdf")
print >>FO, "qual=c(" + ','.join(q_list)  + ')'
print >>FO, "mat=matrix(qual,ncol=%s,byrow=F)" % (read_len)
print >>FO, 'Lab.palette <- colorRampPalette(c("blue", "orange", "red3","red2","red1","red"), space = "rgb",interpolate=c(\'spline\'))'

print >>FO, "heatmap(mat,Rowv=NA,Colv=NA,xlab=\"Position of Read\",ylab=\"Phred Quality Score\",labRow=seq(from="+str(q_min)+",to="+str(q_max)+"),col = Lab.palette(256),scale=\"none\")"
print >>FO, 'dev.off()'
		
		
	